//
// Created by 37417 on 2024/4/5.
//
#include "gtest/gtest.h"
#include "segment_reduction/segment_reduction.h"

TEST(TestUnsortedSegmentSum, test_radix_sort) {
  using namespace segment;
  int size = 10;
  std::vector<int> key_in = {170, 45, 75, 90, 2, 4, 40, 30, 17, 23};
  std::vector<int> keys_out;
  std::vector<int> indices_in = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  std::vector<int> indices_out;

  RadixSort(size, key_in, keys_out, indices_in, indices_out);
}

TEST(TestUnsortedSegmentSum, test_uss_1) {
  //  std::vector<float> data{1, 2, 3, 4, 5, 6, 7, 8, 9};
  float data[9][2]{1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 12, 13, 14, 15, 16, 17, 18, 19};
  float* ptr = &data[0][0];
  std::cout << ptr[14] << std::endl;
  std::vector<int> unsorted_segment_ids{1, 0, 0, 3, 2, 2, 1, 3, 0};
  //  std::vector<float> output(9);
  float output[9][2];
  int input_outer_dim_size = 9;
  int input_inner_dim_size = 2;
  int num_segments = 9;

  using Treduce = float;
  using namespace segment;
  OpKernelContext* ctx;
  using ReductionF = functor::Sum;
  using InitialValueF = functor::Zero<int>;

  std::vector<int> segment_ids_ptr;
  std::vector<int> indices_in = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  std::vector<int> sorted_indices_ptr;

  RadixSort(input_outer_dim_size, unsorted_segment_ids, segment_ids_ptr, indices_in,
            sorted_indices_ptr);

  using T = float;
  using Toffsets = int;
  using Tindices = int;
  using Tsegmentids = int;
  using Tweights = int;
  DisPlay("input", &data[0][0], input_outer_dim_size * input_inner_dim_size);
  SegmentReduceGPU<Treduce, T, Toffsets, Tindices, Tsegmentids, ReductionF, Tweights>(
      ctx, input_outer_dim_size, input_inner_dim_size, num_segments, ReductionF(),
      /*initial_value=*/InitialValueF()(),
      /*empty_segment_value=*/InitialValueF()(), /*is_mean=*/false,
      //      /*is_sqrtn=*/false, /*input=*/data.data(),
      /*is_sqrtn=*/false, /*input=*/&data[0][0],
      /*segment_ids=*/segment_ids_ptr.data(), /*indices=*/sorted_indices_ptr.data(),
      //      /*weights=*/static_cast<int*>(nullptr), output.data());
      /*weights=*/static_cast<int*>(nullptr), &output[0][0]);
  DisPlay("output", &output[0][0], num_segments * input_inner_dim_size);
}
